import matplotlib
# --- 修复：设置非交互式后端 ---
# This must be done BEFORE importing pyplot to prevent GUI-related warnings in a multi-threaded environment like Flask.
matplotlib.use('Agg')

import librosa
import numpy as np
from scipy.signal import butter, filtfilt
from scipy.ndimage import maximum_filter
import matplotlib.pyplot as plt
from typing import Dict, Tuple, Optional
from pathlib import Path

# 导入路径配置
from src.utils import PROCESSOR_RESULTS_PATH

# 为matplotlib设置中文字体
try:
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
except Exception:
    print("Warning: Chinese font 'SimHei' not found, plots may not display Chinese characters correctly.")

class Processor:
    def __init__(self, config: Dict):
        self.sr = config['audio']['sample_rate']
        self.highpass_cutoff = config['audio']['highpass_cutoff']
        self.frame_length = config['stft']['frame_length']
        self.hop_length = config['stft']['hop_length']
        self.peak_energy_threshold = config['peaks']['energy_threshold']
        self.peak_neighborhood_size = config['peaks']['neighborhood_size']
        self.target_peak_density = config['peaks']['target_density']
        self.output_root = PROCESSOR_RESULTS_PATH
        self.output_root.mkdir(parents=True, exist_ok=True)

    def audio_to_peaks(self, audio_path: str) -> np.ndarray:
        try:
            y, _ = librosa.load(audio_path, sr=self.sr, mono=True)
            y_filtered = self._highpass_filter(y)
            log_spectrogram, freqs, times = self._compute_spectrogram(y_filtered)
            peaks = self._find_peaks(log_spectrogram, freqs, times)
            return peaks
        except Exception as e:
            print(f"Error during peak extraction for {audio_path}: {e}")
            return np.array([])

    def process_and_visualize(self, audio_path: str) -> Optional[Dict]:
        try:
            audio_filename = Path(audio_path).stem
            output_dir = self.output_root / audio_filename
            output_dir.mkdir(parents=True, exist_ok=True)
            y, _ = librosa.load(audio_path, sr=self.sr, mono=True)
            y_filtered = self._highpass_filter(y)
            log_spectrogram, freqs, times = self._compute_spectrogram(y_filtered)
            peaks = self._find_peaks(log_spectrogram, freqs, times)
            spec_path = self._save_spectrogram_plot(log_spectrogram, peaks, output_dir / 'spectrogram_with_peaks.png', f'语谱图与峰值点 - {audio_filename}')
            const_path = None
            if peaks.size > 0:
                const_path = self._save_constellation_plot(peaks, output_dir)
            return {"spectrogram_url": Path(spec_path).relative_to(self.output_root).as_posix(), "constellation_url": Path(const_path).relative_to(self.output_root).as_posix() if const_path else None, "peak_count": len(peaks)}
        except Exception as e:
            print(f"Error during visualization for {audio_path}: {e}")
            return None

    def _highpass_filter(self, y: np.ndarray) -> np.ndarray:
        nyquist = 0.5 * self.sr
        cutoff_normalized = self.highpass_cutoff / nyquist
        b, a = butter(5, cutoff_normalized, btype="highpass")
        return filtfilt(b, a, y)

    def _compute_spectrogram(self, y: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        stft_matrix = librosa.stft(y, n_fft=self.frame_length, hop_length=self.hop_length)
        log_spectrogram = librosa.amplitude_to_db(np.abs(stft_matrix), ref=np.max)
        freqs = librosa.fft_frequencies(sr=self.sr, n_fft=self.frame_length)
        times = librosa.frames_to_time(range(log_spectrogram.shape[1]), sr=self.sr, hop_length=self.hop_length)
        return log_spectrogram, freqs, times

    def _find_peaks(self, log_spec: np.ndarray, freqs: np.ndarray, times: np.ndarray) -> np.ndarray:
        footprint = np.ones((self.peak_neighborhood_size, self.peak_neighborhood_size))
        local_max = maximum_filter(log_spec, footprint=footprint, mode='constant', cval=np.min(log_spec))
        detected_peaks_mask = (log_spec == local_max) & (log_spec > self.peak_energy_threshold)
        freq_indices, time_indices = np.nonzero(detected_peaks_mask)
        if freq_indices.size == 0:
            return np.array([])
        all_peak_amps = log_spec[freq_indices, time_indices]
        song_duration = times[-1]
        num_peaks_to_keep = int(self.target_peak_density * song_duration)
        if len(all_peak_amps) <= num_peaks_to_keep:
            peaks = np.array([[times[t], freqs[f]] for f, t in zip(freq_indices, time_indices)])
        else:
            top_peak_indices = np.argsort(all_peak_amps)[::-1][:num_peaks_to_keep]
            top_freq_indices = freq_indices[top_peak_indices]
            top_time_indices = time_indices[top_peak_indices]
            peaks = np.array([[times[t], freqs[f]] for f, t in zip(top_freq_indices, top_time_indices)])
        if peaks.size > 0:
            return peaks[np.argsort(peaks[:, 0])]
        return np.array([])

    def _save_spectrogram_plot(self, log_spec: np.ndarray, peaks: np.ndarray, output_path: Path, title: str) -> str:
        plt.figure(figsize=(15, 8))
        librosa.display.specshow(log_spec, sr=self.sr, hop_length=self.hop_length, x_axis='time', y_axis='log')
        if peaks.size > 0:
            plt.scatter(peaks[:, 0], peaks[:, 1], c='cyan', edgecolor='black', s=50, marker='o', alpha=0.8, linewidth=1)
        plt.colorbar(format='%+2.0f dB')
        plt.title(title)
        plt.savefig(str(output_path), bbox_inches='tight', dpi=150)
        plt.close()
        return str(output_path)

    def _save_constellation_plot(self, peaks: np.ndarray, output_dir: Path) -> str:
        fig, ax = plt.subplots(figsize=(15, 8))
        ax.scatter(peaks[:, 0], peaks[:, 1], s=15, alpha=0.7, c='#00FFFF')
        ax.set_xlabel('时间 (秒)', color='white')
        ax.set_ylabel('频率 (Hz)', color='white')
        ax.set_title(f'星座图 - {output_dir.name}', color='white')
        ax.grid(True, linestyle='--', alpha=0.3)
        ax.set_facecolor('#1f2937')
        fig.set_facecolor('#1f2937')
        ax.tick_params(axis='x', colors='white')
        ax.tick_params(axis='y', colors='white')
        output_path = output_dir / 'constellation_map.png'
        plt.savefig(output_path, bbox_inches='tight', dpi=150, facecolor=fig.get_facecolor())
        plt.close(fig)
        return str(output_path)
